#!/usr/bin/env python3
"""
重新提取音频脚本 - 自定义版本
根据修改后的音频提取逻辑（前后各延长0.5秒），重新从指定视频中提取音频片段并上传到159.138.23.13服务器
"""

import os
import json
import sys
from pathlib import Path
from moviepy.editor import VideoFileClip
import tempfile
import logging
import subprocess

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class AudioReExtractorCustom:
    def __init__(self):
        """初始化音频重新提取器"""
        # 服务器连接配置
        self.server_host = "159.138.23.13"
        self.server_username = "root"
        self.server_password = "UBtest_1"  # 请在这里填入服务器密码
        self.server_base_path = "/data"
        
        if self.server_password == "YOUR_PASSWORD_HERE":
            logging.error("请在代码中设置服务器密码")
            sys.exit(1)
        
    def parse_timestamp(self, timestamp_str):
        """
        解析时间戳字符串为开始和结束时间（毫秒）
        格式: "01:22:14,606 --> 01:22:16,846"
        """
        try:
            start_end = timestamp_str.split(' --> ')
            if len(start_end) != 2:
                raise ValueError(f"无效的时间戳格式: {timestamp_str}")

            start_time = start_end[0].strip()
            end_time = start_end[1].strip()

            def time_str_to_ms(time_str):
                hh_mm_ss, ms = time_str.split(',')
                h, m, s = hh_mm_ss.split(':')
                total_ms = int(h) * 3600 * 1000 + int(m) * 60 * 1000 + int(s) * 1000 + int(ms)
                return total_ms

            return time_str_to_ms(start_time), time_str_to_ms(end_time)
        except Exception as e:
            logging.error(f"解析时间戳失败: {timestamp_str}, 错误: {e}")
            return None, None

    def upload_to_server_scp(self, local_file_path, remote_file_path):
        """
        使用scp命令上传文件到服务器
        """
        try:
            # 确保远程目录存在
            remote_dir = os.path.dirname(remote_file_path)
            mkdir_cmd = f"sshpass -p '{self.server_password}' ssh -o StrictHostKeyChecking=no {self.server_username}@{self.server_host} 'mkdir -p {remote_dir}'"
            subprocess.run(mkdir_cmd, shell=True, check=True, capture_output=True)
            
            # 使用scp上传文件
            scp_cmd = f"sshpass -p '{self.server_password}' scp -o StrictHostKeyChecking=no '{local_file_path}' {self.server_username}@{self.server_host}:'{remote_file_path}'"
            result = subprocess.run(scp_cmd, shell=True, check=True, capture_output=True)
            
            if result.returncode == 0:
                logging.info(f"成功上传文件到服务器: {remote_file_path}")
                return True
            else:
                logging.error(f"上传文件失败: {result.stderr.decode()}")
                return False
                
        except subprocess.CalledProcessError as e:
            logging.error(f"上传文件到服务器失败: {e}")
            return False
        except Exception as e:
            logging.error(f"上传文件到服务器失败: {e}")
            return False

    def extract_audio_segment(self, video_path, start_ms, end_ms, output_path, video_duration):
        """
        从视频中提取音频片段（应用新的延长逻辑）
        """
        try:
            # 应用新的音频提取逻辑：前后各延长0.5秒
            start_sec = start_ms / 1000  # 转换为秒
            start_sec = max(0, start_sec - 0.5)  # 向前延长0.5秒，但不能小于0
            end_sec = end_ms / 1000  # 转换为秒
            end_sec = min(end_sec + 0.5, video_duration)  # 向后延长0.5秒，但不能超过视频时长

            # 确保时间范围有效
            if start_sec >= end_sec:
                logging.error(f"无效的时间范围: {start_sec}s - {end_sec}s")
                return False

            # 加载视频并提取音频片段
            video = VideoFileClip(video_path)
            audio_clip = video.subclip(start_sec, end_sec).audio
            
            # 保存音频片段
            audio_clip.write_audiofile(output_path, codec='mp3', verbose=False, logger=None)
            
            # 清理资源
            audio_clip.close()
            video.close()
            
            logging.info(f"成功提取音频片段: {output_path} ({start_sec:.1f}s - {end_sec:.1f}s)")
            return True
            
        except Exception as e:
            logging.error(f"提取音频片段失败: {e}")
            return False

    def process_dataset(self, ip_dir, video_path):
        """
        处理指定IP目录的数据集
        
        参数:
            ip_dir: IP目录名（如 "IP1fef9dba"）
            video_path: 对应的视频文件路径（如 "data/video/movie.mp4"）
        """
        logging.info(f"开始处理IP目录: {ip_dir}")
        logging.info(f"使用视频文件: {video_path}")
        
        try:
            # 验证视频文件是否存在
            if not Path(video_path).exists():
                logging.error(f"视频文件不存在: {video_path}")
                return False
            
            # 构建数据集文件路径
            dataset_file = Path("output") / ip_dir / f"word_dataset_{ip_dir}_with_style.json"
            
            if not dataset_file.exists():
                logging.error(f"数据集文件不存在: {dataset_file}")
                return False
            
            # 读取JSON文件
            with open(dataset_file, 'r', encoding='utf-8') as f:
                dataset = json.load(f)
            
            # 获取视频时长
            video = VideoFileClip(video_path)
            video_duration = video.duration
            video.close()
            
            logging.info(f"视频文件: {video_path}, 时长: {video_duration:.1f}秒")
            
            # 统计信息
            total_items = 0
            processed_items = 0
            failed_items = 0
            
            # 遍历数据集中的所有单词
            for word, word_data_list in dataset.items():
                for word_data in word_data_list:
                    total_items += 1
                    
                    # 获取必要信息
                    timestamp = word_data.get('timestamp')
                    audio_id = word_data.get('audio_id')
                    audio_path = word_data.get('audio_path')
                    
                    if not all([timestamp, audio_id, audio_path]):
                        logging.warning(f"数据不完整，跳过: {word_data}")
                        failed_items += 1
                        continue
                    
                    # 解析时间戳
                    start_ms, end_ms = self.parse_timestamp(timestamp)
                    if start_ms is None or end_ms is None:
                        logging.warning(f"时间戳解析失败，跳过: {timestamp}")
                        failed_items += 1
                        continue
                    
                    # 创建临时音频文件
                    with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as temp_file:
                        temp_audio_path = temp_file.name
                    
                    try:
                        # 提取音频片段
                        if self.extract_audio_segment(video_path, start_ms, end_ms, temp_audio_path, video_duration):
                            # 构建服务器路径（保持原有路径结构）
                            # 从URL中提取路径: http://159.138.23.13/audio/original/IP1fef9dba/43c7dc48.mp3
                            # 提取: audio/original/IP1fef9dba/43c7dc48.mp3
                            if 'audio/original/' in audio_path:
                                relative_path = audio_path.split('audio/original/')[-1]
                                server_path = f"{self.server_base_path}/audio/original/{relative_path}"
                            else:
                                # 备用方案
                                server_path = f"{self.server_base_path}/audio/original/{ip_dir}/{audio_id}.mp3"
                            
                            # 上传到服务器
                            if self.upload_to_server_scp(temp_audio_path, server_path):
                                logging.info(f"成功重新提取并上传音频: {audio_id} -> {server_path}")
                                processed_items += 1
                            else:
                                logging.error(f"上传音频失败: {audio_id}")
                                failed_items += 1
                        else:
                            logging.error(f"提取音频失败: {audio_id}")
                            failed_items += 1
                            
                    finally:
                        # 清理临时文件
                        if os.path.exists(temp_audio_path):
                            os.unlink(temp_audio_path)
            
            logging.info(f"处理完成 - 总计: {total_items}, 成功: {processed_items}, 失败: {failed_items}")
            return processed_items > 0
            
        except Exception as e:
            logging.error(f"处理数据集失败: {e}")
            return False

def main():
    """主函数"""
    # ========== 在这里配置要处理的IP目录和视频文件 ==========
    
    # 示例配置 - 请根据实际情况修改
    ip_dir = "IP9754cf1e"  # 要处理的IP目录名
    video_path = "data/video/peiqi26-.mp4"  # 对应的视频文件路径（支持.mp4, .mkv等格式）
    
    # 更多示例：
    # ip_dir = "IP1750926257"
    # video_path = "data/video/kungfu_panda.mkv"
    
    # ip_dir = "IP9754cf1e"
    # video_path = "data/video/peppa_pig.mp4"
    
    # =====================================================
    
    logging.info("=== 音频重新提取脚本 ===")
    logging.info(f"配置信息:")
    logging.info(f"  IP目录: {ip_dir}")
    logging.info(f"  视频文件: {video_path}")
    logging.info("=" * 50)
    
    extractor = AudioReExtractorCustom()
    success = extractor.process_dataset(ip_dir, video_path)
    
    if success:
        logging.info("音频重新提取完成！")
        sys.exit(0)
    else:
        logging.error("音频重新提取过程中出现错误")
        sys.exit(1)

if __name__ == "__main__":
    main()
